[None][feat] Enable KVCacheManagerV2 by default for Gemma3 and Gemma4 - #17396
[None][feat] Enable KVCacheManagerV2 by default for Gemma3 and Gemma4#17396erictsai-nv wants to merge 8 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughGemma3 and Gemma4 model variants now enable KV cache manager v2 and select the Python transceiver runtime. Gemma4 multimodal defaults retain the ChangesGemma runtime defaults
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant ModelClass
participant DefaultResolver
participant NIXLRuntime
ModelClass->>DefaultResolver: provide KV cache manager v2 and PYTHON defaults
DefaultResolver->>NIXLRuntime: resolve the Python transceiver runtime
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/modeling_gemma3.py (2)
296-301: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType and document the transceiver runtime hook.
get_preferred_transceiver_runtimehas no parameter or return annotations and no docstring. Match the base hook signature, returnstr, and document why Gemma3 requires the"PYTHON"runtime.As per coding guidelines, annotate every function and use Google-style docstrings for externally usable interfaces.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_gemma3.py` around lines 296 - 301, Type and document the get_preferred_transceiver_runtime classmethod to match the base hook signature, including annotations for cls, pretrained_config, and the str return value. Add a Google-style docstring explaining why Gemma3 requires the "PYTHON" transceiver runtime, while preserving the existing return value.Source: Coding guidelines
288-294: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the required annotations and docstring to this hook.
get_model_defaultshas an untypedllm_argsparameter, a baredictreturn annotation, and no Google-style docstring. Use the type aliases from the base hook and document that this method enables the V2 KV-cache manager.As per coding guidelines, annotate every function and use Google-style docstrings for externally usable interfaces.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_gemma3.py` around lines 288 - 294, Update Gemma3ForCausalLM.get_model_defaults to use the base hook’s type aliases for the llm_args parameter and return value, and add a Google-style docstring describing that it enables the V2 KV-cache manager. Preserve the existing defaults unchanged.Source: Coding guidelines
tensorrt_llm/_torch/models/modeling_gemma4.py (1)
1268-1274: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType and document the transceiver runtime hook.
get_preferred_transceiver_runtimehas no parameter or return annotations and no docstring. Match the base hook signature, returnstr, and document why Gemma4 requires the"PYTHON"runtime.As per coding guidelines, annotate every function and use Google-style docstrings for externally usable interfaces.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_gemma4.py` around lines 1268 - 1274, Update the Gemma4 classmethod get_preferred_transceiver_runtime to match the base hook’s annotated signature, including the pretrained_config parameter type and a str return annotation. Add a Google-style docstring documenting that Gemma4 requires the "PYTHON" transceiver runtime.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4.py`:
- Around line 1263-1266: Update
Gemma4ForConditionalGeneration.get_model_defaults to delegate to
Gemma4ForCausalLM.get_model_defaults(llm_args), preserving the multimodal
wrapper’s existing defaults while inheriting use_kv_cache_manager_v2=True.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma3.py`:
- Around line 296-301: Type and document the get_preferred_transceiver_runtime
classmethod to match the base hook signature, including annotations for cls,
pretrained_config, and the str return value. Add a Google-style docstring
explaining why Gemma3 requires the "PYTHON" transceiver runtime, while
preserving the existing return value.
- Around line 288-294: Update Gemma3ForCausalLM.get_model_defaults to use the
base hook’s type aliases for the llm_args parameter and return value, and add a
Google-style docstring describing that it enables the V2 KV-cache manager.
Preserve the existing defaults unchanged.
In `@tensorrt_llm/_torch/models/modeling_gemma4.py`:
- Around line 1268-1274: Update the Gemma4 classmethod
get_preferred_transceiver_runtime to match the base hook’s annotated signature,
including the pretrained_config parameter type and a str return annotation. Add
a Google-style docstring documenting that Gemma4 requires the "PYTHON"
transceiver runtime.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e6c1195c-2109-49d9-a4a6-8f2b78c7fc8c
📒 Files selected for processing (4)
tensorrt_llm/_torch/models/modeling_gemma3.pytensorrt_llm/_torch/models/modeling_gemma4.pytests/unittest/_torch/modeling/test_modeling_gemma3.pytests/unittest/_torch/modeling/test_modeling_gemma4.py
5b8bd57 to
5bfdca3
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #64567 [ run ] triggered by Bot. Commit: |
|
PR_Github #64567 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The mechanism and the [None] tag are fine for a default flip, but the coverage doesn't match what the description claims.
Model defaults are keyed by the checkpoint's architectures[0] (model_loader.py:415, :452). Gemma3 4B/12B/27B ship as Gemma3ForConditionalGeneration, which resolves to Gemma3VLM (modeling_gemma3vl.py:178) — that class has neither hook, so it keeps V1. The description says "all Gemma3 variants" and lists 27B as validated on the v2cpp arm; either the validation used a text-only checkpoint, or the resolution went somewhere I'm not seeing. Please reconcile — add the hook to Gemma3VLM if VLM V2 was actually what you validated, otherwise narrow the description to Gemma3ForCausalLM.
docs/source/features/kvcache.md:117 currently documents only hybrid Mamba as selecting V2 under auto. Since this changes OOTB behavior for two shipped model families, worth a sentence there.
Also worth confirming: [None] is right for a default flip that changes runtime behavior for released models — if there's a Gemma-V2-enablement JIRA, cite it so the rollback story is traceable.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modeling/test_modeling_gemma3.py`:
- Around line 698-740: Register
tests/unittest/_torch/modeling/test_modeling_gemma3.py in the appropriate
tests/integration/test_lists/test-db/ and qa/ test lists, matching the existing
registration for test_modeling_gemma4.py in
tests/integration/test_lists/test-db/l0_b200.yml. Ensure the Gemma3 tests
test_gemma3_defaults_resolve_to_v2 and test_gemma3_explicit_user_setting_wins,
along with the two Gemma4 tests in
tests/unittest/_torch/modeling/test_modeling_gemma4.py, are executed by the
registered suites; run pytest tests/unittest/ and report the resulting coverage.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c453d524-b420-4505-a2a1-69bb706d9a24
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_gemma3.pytensorrt_llm/_torch/models/modeling_gemma3vl.pytensorrt_llm/_torch/models/modeling_gemma4.pytensorrt_llm/_torch/models/modeling_gemma4mm.pytests/unittest/_torch/modeling/test_modeling_gemma3.pytests/unittest/_torch/modeling/test_modeling_gemma4.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/models/modeling_gemma3.py
|
Docs: deferring to #16942, which adds a "Selecting the KV Cache Manager" table — it'll be merged soon; I'll rebase and add a Gemma row there. No Gemma-V2 enablement ticket exists; keeping [None]. |
ff21cad to
a7045f9
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/source/features/kvcache.md (1)
1-1: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the required NVIDIA copyright header.
docs/source/features/kvcache.mdis modified, but it starts with# KV Cache Systemand has no NVIDIA copyright header. Add the repository-standard header at the top with year2026.As per coding guidelines, modified files must carry the NVIDIA copyright header with the year of the latest meaningful modification.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/source/features/kvcache.md` at line 1, Add the repository-standard NVIDIA copyright header with year 2026 at the beginning of the kvcache.md document, before the existing “KV Cache System” heading.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/source/features/kvcache.md`:
- Line 91: Update the Gemma3 / Gemma4 row in the KV-cache documentation table to
explicitly scope the entry to non-hybrid text and multimodal models, or split it
into separate rows reflecting the routing behavior; do not imply that hybrid
Gemma4 models use the documented model-default path.
In `@tests/unittest/_torch/modeling/test_modeling_gemma4.py`:
- Around line 4054-4081: Restrict the NIXL/transceiver-runtime assertions in
test_gemma4_defaults_resolve_to_v2 to Gemma4ForCausalLM, since
Gemma4ForConditionalGeneration rejects multimodal disaggregated inference. Add a
separate Gemma4ForConditionalGeneration test that resolves and verifies the V2
KV-cache-manager and FLASHINFER defaults without configuring
cache_transceiver_config.
---
Outside diff comments:
In `@docs/source/features/kvcache.md`:
- Line 1: Add the repository-standard NVIDIA copyright header with year 2026 at
the beginning of the kvcache.md document, before the existing “KV Cache System”
heading.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e161493c-686e-4498-8e98-75d7adb6eac4
📒 Files selected for processing (7)
docs/source/features/kvcache.mdtensorrt_llm/_torch/models/modeling_gemma3.pytensorrt_llm/_torch/models/modeling_gemma3vl.pytensorrt_llm/_torch/models/modeling_gemma4.pytensorrt_llm/_torch/models/modeling_gemma4mm.pytests/unittest/_torch/modeling/test_modeling_gemma3.pytests/unittest/_torch/modeling/test_modeling_gemma4.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/_torch/models/modeling_gemma3vl.py
- tensorrt_llm/_torch/models/modeling_gemma4.py
- tensorrt_llm/_torch/models/modeling_gemma4mm.py
- tensorrt_llm/_torch/models/modeling_gemma3.py
|
Now we use get preferred kvcm like transceiver. Please refer to #16060 |
a7045f9 to
37c5ec3
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Re-checked the three earlier concerns against the current head; all three are genuinely fixed, not just marked resolved:
Gemma3ForConditionalGeneration→Gemma3VLMnow carries both hooks ([modeling_gemma3vl.py:180](https://github.com/NVIDIA/TensorRT-LLM/pull/17396/files#diff-53f0a764a7ee3ee4eef673b081e3bcc3c0d545e856882db2414524a7929113a1R180),:189), so the 4B/12B/27B checkpoints actually pick up V2 —get_registered_model_class("Gemma3ForConditionalGeneration")is asserted in the new test, which is the property that was wrong before.Gemma4MultimodalModelBasedeclaresget_preferred_transceiver_runtime() -> "PYTHON"([modeling_gemma4mm.py:571](https://github.com/NVIDIA/TensorRT-LLM/pull/17396/files#diff-c199d40bc9ba6c34e1a2ed072172254df5e14e0d33790fae775615591b4a744fR571)), so the NIXL route no longer silently demotes MM back to V1 in_resolve_kv_cache_manager_v2_auto.Gemma4UnifiedForConditionalGenerationinherits it (modeling_gemma4_unified.py:224).- The unit tests now replay the production ordering from
model_loader.load_config_and_apply_defaults(defaults merge →_resolve_transceiver_runtime_auto→_resolve_kv_cache_manager_v2_auto) instead of re-asserting a dict literal, and the explicit-user-value tests cover both directions.
The V2 preference having moved from get_model_defaults into get_preferred_kv_cache_manager_version (37c5ec3) is the right call — a kv_cache_config entry in the defaults dict would have bypassed the compatibility demotions (CPP transceiver, two-model spec dec) that the "auto" resolver applies.
Two non-blocking notes:
Gemma4UnifiedForConditionalGenerationinherits the preference but is unvalidated, as the description says. Fine to ship given it shares the text backbone, but it's the one arch where a surprise would show up first.- Nothing under
tests/integration/test_lists/test-db/changed; the coverage claim rests on existing Gemma tests now taking the"auto"path. That's reasonable, just worth being explicit that no new integration case guards the disagg/NIXL combination the transceiver hook exists for.
LGTM.
|
/bot run --disable-fail-fast |
|
PR_Github #65456 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with nits.
|
PR_Github #65456 [ run ] completed with state
|
920b745 to
5a1da50
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65550 [ run ] triggered by Bot. Commit: |
|
PR_Github #65550 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65735 [ run ] triggered by Bot. Commit: |
5a1da50 to
0c1dd77
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65745 [ run ] triggered by Bot. Commit: |
|
PR_Github #65735 [ run ] completed with state |
|
PR_Github #65745 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65800 [ run ] triggered by Bot. Commit: |
|
PR_Github #65800 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65839 [ run ] triggered by Bot. Commit: |
|
PR_Github #65839 [ run ] completed with state
|
Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…t_model_defaults Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…solution-path tests Signed-off-by: Eric Tsai <ertsai@nvidia.com>
Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…ager_version Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…-M3; add NIXL V2-retention test Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…VIDIA#16787) Signed-off-by: Eric Tsai <ertsai@nvidia.com>
…made it prefer the Python transceiver) Signed-off-by: Eric Tsai <ertsai@nvidia.com>
0c1dd77 to
a5c6e94
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #66117 [ run ] triggered by Bot. Commit: |
|
PR_Github #66117 [ run ] completed with state
|
5 participants
Dev Engineer Review
KVCacheManagerV2by default.FLASHINFERbackend default."PYTHON"transceiver runtime for disaggregated NIXL serving.QA Engineer Review
test_gemma3_defaults_resolve_to_v2.test_gemma3_explicit_user_setting_wins.test_gemma4_defaults_resolve_to_v2.test_gemma4_explicit_user_setting_wins.test-db/orqa/.Description
Enables
KVCacheManagerV2by default for Gemma3 and Gemma4.All four Gemma entry points —
Gemma3ForCausalLM,Gemma3VLM(Gemma3ForConditionalGenerationcheckpoints, i.e. 4B/12B/27B),Gemma4ForCausalLM, andGemma4MultimodalModelBase— declareget_preferred_kv_cache_manager_version() -> "V2"(the per-model preference hook introduced in #16060).Gemma3VLMand the Gemma4 entry points also declareget_preferred_transceiver_runtime() -> "PYTHON"so that disaggregated serving over NIXL retains V2 instead of falling back to V1 (Gemma3ForCausalLMgained the same hook upstream in #16787). The preference is adopted only when the user leaveskv_cache_config.use_kv_cache_manager_v2at"auto"; an explicit user value always wins.Notes:
_non_hybrid_kv_cache_manager_cls(_util.py); this change extends the default to non-hybrid Gemma4 and all Gemma3 variants.Gemma4UnifiedForConditionalGeneration(12B) is hybrid attention, so it is routed to V2 unconditionally as well — the preference it inherits throughGemma4MultimodalModelBaseis a no-op for it. Validated on the real 12B checkpoint regardless: text + image outputs are sane, and the cpp/python V2 backends produce bitwise-identical tokens (B200, greedy).get_preferred_transceiver_runtime() -> "PYTHON"to DeepSeek-V4 and MiniMax-M3 (plusget_preferred_kv_cache_manager_version() -> "V2"for M3): both run V2-core cache managers that the C++ transceiver cannot drive, and their disaggregated tests already pin NIXL + PYTHON — the fully-"auto" route now resolves to the same combination. MiniMax-M2 is deliberately left without a preference: its disaggregated serving is unvalidated, and the absence keeps that visible.Since #14047 landed,
TLLM_KV_CACHE_MANAGER_V2_BACKENDdefaults tocpp, so this PR moves Gemma onto the C++ V2 core.Validation
H100 (agg + disagg) and B200 (NVFP4 MMMU): Gemma3 1B/27B (27B = VLM entry point) and Gemma4 26B-A4B (multimodal entry point) all passed across v1/v2py/v2cpp arms, accuracy within run noise. Gemma4 unified 12B: real-checkpoint smoke on B200, v2py/v2cpp bitwise identical with sane multimodal outputs.
Test Coverage
test_registered_models_prefer_v2(extended): the five Gemma architectures and the two MiniMax-M3 architectures join the V2-preference list.test_registered_models_keep_v2_on_nixl(new, next to the list above): for every model preferring both V2 and the Python transceiver, resolve the fully-"auto" NIXL route in production order (transceiver runtime first, then KV cache manager) and assert V2 + PYTHON survive. MiniMax-M2 (disagg unvalidated) is documented exclusion.TestKvCacheManagerV2AutoResolution.Existing Gemma3/4 integration tests that do not set
use_kv_cache_manager_v2explicitly now run through the"auto"resolution path and exercise V2 by default.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.